import csv
import shutil
import sys

import numpy as np
import orekit
import pandas as pd
from orekit.pyhelpers import setup_orekit_curdir
from org.hipparchus.geometry.euclidean.threed import Vector3D
from org.orekit.data import DataProvidersManager, DirectoryCrawler
from org.orekit.orbits import KeplerianOrbit
from org.orekit.propagation.analytical.tle import TLE, TLEPropagator
from org.orekit.time import AbsoluteDate, DateTimeComponents, TimeScalesFactory, Month
from org.orekit.frames import FramesFactory
from org.orekit.utils import IERSConventions
from datetime import datetime, timedelta
from sgp4.api import Satrec, WGS84
from orbitP.script import config
from tqdm import tqdm
import os


# 初始化 Orekit 环境
orekit.initVM()
setup_orekit_curdir(config.orekitCurPath)

# TLE数据
tle1 = ""
tle2 = ""
beginDate = datetime.strptime("2023-01-01","%Y-%m-%d")
endDate = datetime.strptime("2025-01-01","%Y-%m-%d")
dataLastTLEPath = "../dataset/dataLastTLE.txt"
dataTLEPath = "../dataset/dataTLE.txt"
dataPOEORBDir = "../dataset/dataPOEORB"
dataOrekitDir = "../dataset/dataOrekit/"

dataDict = {}
dateRes = {}

def get_file_paths(directory):
    file_paths = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            file_path = os.path.join(root, file)
            file_paths.append(os.path.normpath(file_path))
    return file_paths

# 将TLE的时间部分转换为标准的datetime对象
def epoch2datetime(year, day):
    """
    将 TLE 中的 epoch year 和 epoch day 转换为标准 datetime 格式。
    :param year: TLE 中的年份，例如 24 表示 2024 年
    :param day: TLE 中的天数，格式为小数，表示该年的第几天
    :return: 对应的 datetime 对象 (UTC)
    """
    if year < 57:
        year = 2000 + year
    else:
        year = 1900 + year
    epoch_date = datetime(year, 1, 1) + timedelta(days=day - 1)
    return epoch_date

# 计算卫星在 TLE 时刻的位置和速度
def getTLEMsg(line1, line2):
    satellite = Satrec.twoline2rv(line1, line2, WGS84)
    # 提取TLE的信息
    year = satellite.epochyr
    day = satellite.epochdays
    time = epoch2datetime(year, day)

    # 将历元转为儒略日
    jd, fr = satellite.jdsatepoch, satellite.jdsatepochF

    # 使用SGP4模型计算位置和速度
    e, p, v = satellite.sgp4(jd, fr)
    if e != 0:
        raise ValueError(f"SGP4 Propagation error, code: {e}")

    return time, p, v

def getMonth(month):
    # 将月份数字转换为Month枚举
    months = {
        1: Month.JANUARY,
        2: Month.FEBRUARY,
        3: Month.MARCH,
        4: Month.APRIL,
        5: Month.MAY,
        6: Month.JUNE,
        7: Month.JULY,
        8: Month.AUGUST,
        9: Month.SEPTEMBER,
        10: Month.OCTOBER,
        11: Month.NOVEMBER,
        12: Month.DECEMBER
    }
    return months[month]
def getPredictMsg_itrf(line1, line2, time_utc):
    tle = TLE(line1, line2)
    utc_time = datetime.strptime(time_utc, "%Y-%m-%d %H:%M:%S")
    # 创建 TLEPropagator
    propagator = TLEPropagator.selectExtrapolator(tle)

    # 设置目标时间（2024年1月1日0时0分2秒）
    utc = TimeScalesFactory.getUTC()
    month_enum = getMonth(utc_time.month)

    target_time = AbsoluteDate(
        utc_time.year, utc_time.month, utc_time.day,
        utc_time.hour, utc_time.minute,
        utc_time.second + utc_time.microsecond / 1e6,
        utc
    )

    # 传播到目标时间并获取轨道状态
    orbitProp = propagator.propagate(target_time)

    # 获取ITRF坐标系
    itrf = FramesFactory.getITRF(IERSConventions.IERS_2010, False)

    # 获取ITRF坐标系中的PV坐标
    pv_coordinates_itrf = orbitProp.getPVCoordinates(itrf)

    kepler = KeplerianOrbit(orbitProp.getOrbit())


    # 获取位置、速度和加速度
    p_itrf = pv_coordinates_itrf.getPosition()
    v_itrf = pv_coordinates_itrf.getVelocity()
    a_itrf = pv_coordinates_itrf.getAcceleration()

    # 将 p_rsw 和 v_rsw 转换为列表
    p_itrf_list = [p_itrf.getX(), p_itrf.getY(), p_itrf.getZ()]
    v_itrf_list = [v_itrf.getX(), v_itrf.getY(), v_itrf.getZ()]
    a_itrf_list = [a_itrf.getX(), a_itrf.getY(), a_itrf.getZ()]

    # 提取六根数
    a = kepler.getA()  # 半长轴（m）
    e = kepler.getE()  # 偏心率
    i = kepler.getI()  # 轨道倾角（rad）
    raan = kepler.getRightAscensionOfAscendingNode()  # 升交点赤经（rad）
    omega = kepler.getPerigeeArgument()  # 近地点幅角（rad）
    mean_anomaly = kepler.getMeanAnomaly()  # 平近点角（rad）

    # 返回一个合并的列表
    return p_itrf_list, v_itrf_list, a_itrf_list,[a,e,i,raan,omega,mean_anomaly]


def getPredictMsg_rsw(line1, line2, time_utc):
    tle = TLE(line1, line2)
    utc_time = datetime.strptime(time_utc, "%Y-%m-%d %H:%M:%S")

    propagator = TLEPropagator.selectExtrapolator(tle)
    utc = TimeScalesFactory.getUTC()
    target_time = AbsoluteDate(
        utc_time.year, utc_time.month, utc_time.day,
        utc_time.hour, utc_time.minute,
        utc_time.second + utc_time.microsecond / 1e6,
        utc
    )

    # 在惯性系（如 GCRF）中传播
    inertial_frame = FramesFactory.getGCRF()
    orbit_inertial = propagator.propagate(target_time).getPVCoordinates(inertial_frame)

    # 提取惯性系下的位置、速度、加速度
    p_inertial = orbit_inertial.getPosition()
    v_inertial = orbit_inertial.getVelocity()
    a_inertial = orbit_inertial.getAcceleration()

    # 计算 RSW 基向量
    # 1. 径向方向 R (手动归一化)
    p_norm = p_inertial.getNorm()
    r_hat = Vector3D(p_inertial.getX() / p_norm,
                     p_inertial.getY() / p_norm,
                     p_inertial.getZ() / p_norm)

    # 2. 法向 W (先计算角动量，再手动归一化)
    h = p_inertial.crossProduct(v_inertial)
    h_norm = h.getNorm()
    w_hat = Vector3D(h.getX() / h_norm,
                     h.getY() / h_norm,
                     h.getZ() / h_norm)

    # 3. 切向方向 S
    s_hat = w_hat.crossProduct(r_hat)

    # 转换函数
    def to_rsw(vec):
        return [vec.dotProduct(r_hat), vec.dotProduct(s_hat), vec.dotProduct(w_hat)]

    p_rsw = to_rsw(p_inertial)
    v_rsw = to_rsw(v_inertial)
    a_rsw = to_rsw(a_inertial)

    return p_rsw, v_rsw, a_rsw

def getCloseTLE(dateTime):
    closeTLE = []
    closeDelta = None
    with open(dataTLEPath, "r", encoding="utf-8") as file:
        lines = file.readlines()
        for i in range(0, len(lines), 2):
            TLE_arr = lines[i:i + 2]
            tle1 = TLE_arr[0].strip()
            tle2 = TLE_arr[1].strip()
            time_TLE, p_TLE, v_TLE = getTLEMsg(tle1, tle2)
            nowDelta = abs(time_TLE-dateTime)
            if closeDelta == None:
                closeDelta = nowDelta
                closeTLE = [tle1,tle2]
            elif closeDelta > nowDelta:
                closeDelta = nowDelta
                closeTLE = [tle1,tle2]
    return closeTLE

def setIn(dateRes,dateTime,p_itrf,v_itrf,a_itrf):
    dateRes['date'].append(dateTime)
    dateRes['X'].append(p_itrf.getX())
    dateRes['Y'].append(p_itrf.getY())
    dateRes['Z'].append(p_itrf.getZ())
    dateRes['VX'].append(v_itrf.getX())
    dateRes['VY'].append(v_itrf.getY())
    dateRes['VZ'].append(v_itrf.getZ())
    dateRes['AX'].append(a_itrf.getX())
    dateRes['AY'].append(a_itrf.getY())
    dateRes['AZ'].append(a_itrf.getZ())

if __name__ == "__main__":
    if os.path.exists(dataOrekitDir):
        shutil.rmtree(dataOrekitDir)
    if not os.path.exists(dataOrekitDir):
        os.mkdir(dataOrekitDir)
    with open(dataLastTLEPath, "r", encoding="utf-8") as file:
        lines = file.readlines()
        for i in range(0, len(lines), 2):
            TLE_arr = lines[i:i + 2]
            tle1 = TLE_arr[0].strip()
            tle2 = TLE_arr[1].strip()
            time_TLE, p_TLE, v_TLE = getTLEMsg(tle1, tle2)
            dataDict[time_TLE.date().strftime("%Y-%m-%d")]=[tle1,tle2]

    filesPath = get_file_paths(dataPOEORBDir)
    t_bar = tqdm(filesPath, total=len(filesPath))
    for filePath in t_bar:
        fileName = os.path.basename(filePath).replace(".csv", "")
        fileDate = datetime.strptime(fileName, "%Y-%m-%d")
        if fileDate < beginDate or fileDate > endDate:
            continue
        t_bar.set_description(f"{fileName}")
        writePath = dataOrekitDir + os.path.basename(filePath)

        with open(filePath, "r", encoding="utf-8") as file:
            csv_reader = csv.reader(file)
            next(csv_reader)
            for row in csv_reader:
                dateTime = datetime.strptime(row[0], "%Y-%m-%d %H:%M:%S")
                date = dateTime.strftime('%Y-%m-%d')
                time_utc_pre = datetime.strptime(row[0], "%Y-%m-%d %H:%M:%S") - timedelta(
                    days=1)  # 使用前一天最晚的TLE预测
                nowTLE = dataDict[time_utc_pre.date().strftime("%Y-%m-%d")]
                # nowTLE = getCloseTLE(dateTime)
                p_itrf,v_itrf,a_itrf, kepler_data = getPredictMsg_itrf(nowTLE[0], nowTLE[1], row[0])
                # setIn(dateRes,dateTime,p_itrf,v_itrf,a_itrf)
                # p_rsw,v_rsw,a_rsw = getPredictMsg_rsw(nowTLE[0], nowTLE[1], row[0])
                # dateNow = p_itrf + v_itrf + a_itrf + kepler_data
                dateNow = p_itrf + v_itrf
                dateNow.insert(0, dateTime)

                if date in dateRes:
                    dateRes[date].append(dateNow)
                else:
                    dateRes[date] = []
                    dateRes[date].append(dateNow)

    for date, data in dateRes.items():
        writePath = f"../dataset/dataOrekit/{date}.csv"
        # df = pd.DataFrame(data, columns=["date", "X", "Y", "Z", "VX", "VY", "VZ", "AX", "AY", "AZ","a", "e","i", "raan", "omega", "mean_anomaly"])
        df = pd.DataFrame(data, columns=["date", "X", "Y", "Z", "VX", "VY", "VZ"])
        df['date'] = pd.to_datetime(df['date'])
        df.to_csv(writePath, index=False)